import paramiko
import time
import os
import stat
from tqdm import tqdm


class SSHConnection(object):
    def __init__(self, adress='school') -> None:
        self.ssh = paramiko.SSHClient()
        self.adress = adress
        self.ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        if adress == 'school':
            self.ssh.connect('10.212.70.131', username='manxk',
                             password='xkMangroup')
            self.server_path = '/gs/home/manxk/baixj/'
        elif adress == 'lvph':
            self.ssh.connect('172.17.176.248',
                             username='yizq', password='2219145')
            self.server_path = '/home/yizq/MD/lammps_run_space/'
        self.shell = self.ssh.invoke_shell()
        self.sftp = self.ssh.open_sftp()

    def std_out(self, command='', output=''):
        list = output.split('\n')
        n = 0
        for i in list:
            if n == 0 and command in i:
                n = 1
                print(i)
                continue
            if n == 0:
                continue
            print(i)

    def task_chack(self, name='MD-shear'):
        stdin, stdout, stderr = self.ssh.exec_command(f'squeue -n {name}\n')
        list = stdout.readlines()
        k = 0
        for i in range(len(list)-1):
            if f'JOBID' in list[i]:
                k = 1
                if list[i+1].find(name) != -1:
                    self.finaltime = list[i+1].split()[5]
                    return 1
                else:
                    return 0
            if k == 0:
                continue
        return 0

    def auto_cmd(self, command):
        stdin, stdout, stderr = self.ssh.exec_command(command)
        for i in stdout.readlines():
            print(i.replace('\n', ''))
        # print(stdout.readlines())

    def key_cmd(self):
        while True:
            command = input('command:')
            if command == 'exit':
                break
            self.shell.send(command+'\n')
            time.sleep(1)
            output = self.shell.recv(1024000).decode('utf-8')
            # print(output)
            self.std_out(command, output)

    def mission_sub(self, mission_space_path):
        """
        Execute the 'bash run' command on remote directories specified by mission_space_path.
        Args:
            mission_space_path (list): List of paths on the remote server.
        """
        for remote_dir in mission_space_path:
            try:
                # Navigate to the remote directory and execute the command
                remote_dir = '~' + \
                    remote_dir[remote_dir.index('manxk')+len('manxk'):]

                cmd = f"cd {remote_dir} && pwd  && bash run"
                stdin, stdout, stderr = self.ssh.exec_command(cmd)

                # Collect the output and errors (optional, for debugging)
                stdout_output = stdout.read().decode()
                stderr_output = stderr.read().decode()

                if stdout_output:
                    print(f"Output from {remote_dir}:\n{stdout_output}")
                if stderr_output:
                    print(f"Error from {remote_dir}:\n{stderr_output}")

            except Exception as e:
                print(f"Failed to execute in {remote_dir}: {e}")

    def close(self):
        # 关闭SFTP连接
        self.sftp.close()
        # 关闭SSH连接
        self.ssh.close()

    def upload(self, local_path, remote_path):
        # 上传文件
        self.sftp.put(local_path, remote_path)

    def download(self, remote_path, local_path):
        # 下载文件
        self.sftp.get(remote_path, local_path)

    def download_dir_yiziqiang(self, remote_dir, local_dir):
        # 下载目录
        def chack(local_dir, n=1):
            if not os.path.exists(local_dir) and not os.path.exists(local_dir+'.txt'):
                os.makedirs(local_dir)
                return local_dir
            else:
                local_dir = local_dir[:-1]+str(n)
                return chack(local_dir, n+1)
        sftp = self.sftp
        local_dir = chack(local_dir+'_N0')
        # print('downoloaddir=',local_dir)
        for filename in sftp.listdir(remote_dir):
            remote_path = remote_dir + '/' + filename
            local_path = local_dir + '/' + filename
            if stat.S_ISDIR(sftp.stat(remote_path).st_mode):
                self.download_dir(remote_path, local_path)
            else:
                sftp.get(remote_path, local_path)
        return local_dir

    from tqdm import tqdm
    import os
    import stat

    def download_dir(self, remote_dir, local_dir, pbar=None):
        """
        将远程目录 remote_dir 复制到本地目录 local_dir，保留结构。
        """
        def count_items_in_dir(remote_dir):
            """
            统计远程目录 remote_dir 中的所有文件和目录项数（包括子目录的内容）。
            """
            sftp = self.sftp
            total_items = 0
            remote_items = sftp.listdir(remote_dir)

            for filename in remote_items:
                remote_path = f"{remote_dir}/{filename}"
                if stat.S_ISDIR(sftp.stat(remote_path).st_mode):
                    total_items += count_items_in_dir(remote_path)  # 递归统计子目录
                else:
                    total_items += 1

            return total_items

        def find_next_suffix(parent_dir, base_name, suffix='N'):
            import re
            """
            查找以 base_name + suffix + 数字 结尾的文件或目录，并返回下一个可用的名称。

            参数：
                parent_dir (str): 父目录的路径。
                base_name (str): 基础名称。
                suffix (str): 后缀标识，默认为 'N'。

            返回：
                str: 新的目录或文件名称。
            """
            # 定义匹配以 base_name + suffix + 数字 结尾的正则表达式
            pattern = re.compile(re.escape(base_name) +
                                 re.escape(suffix) + r'(\d+)$')

            # 初始化最大数字为 -1，表示尚未找到匹配项
            max_num = -1

            try:
                # 遍历父目录中的所有文件和子目录
                for name in os.listdir(parent_dir):
                    match = pattern.match(name)
                    if match:
                        num = int(match.group(1))
                        if num > max_num:
                            max_num = num
            except FileNotFoundError:
                print(f"父目录 '{parent_dir}' 不存在。")
            except PermissionError:
                print(f"没有权限访问父目录 '{parent_dir}'。")
            except Exception as e:
                print(f"遍历目录时出错: {e}")

            # 根据找到的最大数字确定新的目录名称
            if max_num == -1:
                new_name = f"{base_name}{suffix}0"
            else:
                new_name = f"{base_name}{suffix}{max_num + 1}"

            return new_name

        sftp = self.sftp

        # 提取 remote_dir 的最后一级目录名
        remote_dir_name = os.path.basename(remote_dir.rstrip('/'))
        target_local_dir = os.path.join(local_dir, remote_dir_name)

        # 定义目标本地目录路径
        if os.path.exists(target_local_dir):
            print(f"本地目录 '{target_local_dir}' 已存在")

            # 获取目标目录的父目录和基名
            parent_dir, base_name = os.path.split(
                os.path.abspath(target_local_dir))

            # 获取新的目录名称
            new_dir_name = find_next_suffix(parent_dir, base_name, suffix='N')
            new_dir_path = os.path.join(parent_dir, new_dir_name)

            print(f"创建新的目录: {new_dir_path}")

            try:
                os.makedirs(new_dir_path)
                # print(f"目录 '{new_dir_path}' 创建成功")
            except Exception as e:
                print(f"创建目录时出错: {e}")
            target_local_dir = new_dir_path
        else:
            try:
                # 创建目标本地目录
                os.makedirs(target_local_dir)
                # print(f"已创建目录: {target_local_dir}")
            except Exception as e:
                print(f"创建目录时出错: {e}")

        # 获取远程目录内容
        remote_items = sftp.listdir(remote_dir)

        # 如果进度条不存在，初始化，并设置标志
        own_pbar = False
        if pbar is None:
            total_items = count_items_in_dir(remote_dir)  # 统计总数
            file_name = remote_dir.split('/')[-1]
            pbar = tqdm(total=total_items,
                        desc=f"Downloading {file_name}", unit="item")
            own_pbar = True  # 标记为自身创建的进度条

        for filename in remote_items:
            remote_path = f"{remote_dir}/{filename}"
            local_path = os.path.join(target_local_dir, filename)

            # 判断远程路径是文件还是目录
            if stat.S_ISDIR(sftp.stat(remote_path).st_mode):
                # 如果是目录，递归调用，并传递进度条
                self.download_dir(remote_path, target_local_dir, pbar)
            else:
                # 如果是文件，下载文件
                sftp.get(remote_path, local_path)
                pbar.update(1)  # 更新进度条

        # 如果是自身创建的进度条，关闭它
        if own_pbar:
            pbar.close()

    def upload_folder(self, local_folder, remote_folder):
        # 确保远程目录存在，同本部
        # self._make_remote_dir(remote_folder)

        # 获取目标远程路径，将整个 local_folder 作为子目录移动到 remote_folder
        folder_name = os.path.basename(
            local_folder.rstrip('/\\'))  # 获取 local_folder 的名称
        remote_target_folder = os.path.join(
            remote_folder, folder_name).replace('\\', '/')

        # # 确保远程目标目录存在
        # self._ensure_remote_dir(remote_target_folder)

        # 遍历 local_folder 并上传内容
        for root, dirs, files in os.walk(local_folder):
            # 计算相对路径
            relative_path = os.path.relpath(root, local_folder)
            # 构建远程路径
            remote_path = os.path.join(
                remote_target_folder, relative_path).replace('\\', '/')
            # 确保远程子目录存在
            self._ensure_remote_dir(remote_path)

            # 上传文件
            for file in files:
                local_file = os.path.join(root, file)
                remote_file = os.path.join(
                    remote_path, file).replace('\\', '/')
                self.sftp.put(local_file, remote_file)

        print(f'上传文件：{local_folder} 到 {remote_folder}')

    def _make_remote_dir(self, remote_directory):
        """创建目录"""
        try:
            self.sftp.stat(remote_directory)
        except FileNotFoundError:
            # 递归创建目录
            parent_dir = os.path.dirname(remote_directory)
            self._make_remote_dir(parent_dir)
            self.sftp.mkdir(remote_directory)
            # print(f'创建远程目录：{remote_directory}')

    def _ensure_remote_dir(self, remote_directory):
        """确保远程目录存在，不存在则创建，并检查是否为空"""
        try:
            # 尝试访问目录
            self.sftp.stat(remote_directory)
        except FileNotFoundError:
            # 如果目录不存在，递归创建
            parent_dir = os.path.dirname(remote_directory)
            self._make_remote_dir(parent_dir)
            self.sftp.mkdir(remote_directory)
            # print(f'创建远程目录：{remote_directory}')
        else:
            # 检查目录是否为空
            is_empty = self._is_remote_dir_empty(remote_directory)
            if not is_empty:
                user_input = input(f"目录 {remote_directory} 不为空。是否清空？(y/n): ")
                if user_input.lower() == 'y':
                    # 双重确认
                    confirm_input = input(
                        f"再次确认是否清空目录 {remote_directory}？(y/n): ")
                    if confirm_input.lower() == 'y':
                        self._clear_remote_directory(remote_directory)
                        print(f"清空远程目录：{remote_directory}")
                    else:
                        print("已取消清空操作")

    def _is_remote_dir_empty(self, remote_directory):
        """检查远程目录是否为空"""
        try:
            files = self.sftp.listdir(remote_directory)
            return len(files) == 0
        except Exception as e:
            print(f"检查目录是否为空时出错：{e}")
            return False

    def _clear_remote_directory(self, remote_directory):
        """清空远程目录"""
        try:
            for item in self.sftp.listdir_attr(remote_directory):
                item_path = os.path.join(remote_directory, item.filename)
                if stat.S_ISDIR(item.st_mode):  # 如果是目录
                    self._clear_remote_directory(item_path)  # 递归清空子目录
                    self.sftp.rmdir(item_path)
                else:  # 如果是文件
                    self.sftp.remove(item_path)

        except Exception as e:
            print(f"清空目录时出错：{e}")

    def exec_command(self, command):
        # 执行多个命令
        if command[-1:] != '\n':
            self.shell.send(command+'\n')
        else:
            self.shell.send(command)
        time.sleep(5)
        output = self.shell.recv(1024000).decode('utf-8')
        self.std_out(command, output)
        # self.shell.send('clear \n')


if __name__ == "__main__":
    # host=input('host:')
    p = SSHConnection('school')
    p.auto_cmd('squeue')
    p.close()
